import asyncio
import io
from fastmcp import FastMCP
from fastmcp import Client
from datetime import datetime
import pandas as pd
from typing import List, Literal, Any
import yfinance as yf
from concurrent.futures import ThreadPoolExecutor
import talib

# 这几个工具都使用到yfinance接口，测试时获取不到数据

mcp = FastMCP("demo.mcp")

# 获取指定股票代码的综合数据，包括基本指标、财报日历、新闻头条、分析师评级和升级/降级记录
@mcp.tool()
def get_ticker_data(
    ticker: str,
    max_news: int = 5,
    max_recommendations: int = 5,
    max_upgrades: int = 5
) -> dict[str, Any]:
    """
    Get comprehensive stock ticker data including basic metrics, earnings calendar,
    news headlines, analyst recommendations, and upgrades/downgrades.

    This function uses the `yfinance` library to fetch data from Yahoo Finance.
    It runs multiple requests in parallel to improve performance.

    Args:
        ticker (str): The stock ticker symbol (e.g., "AAPL").
        max_news (int, optional): Max number of news items to return. Default: 5.
        max_recommendations (int, optional): Max number of recommendation rows. Default: 5.
        max_upgrades (int, optional): Max number of upgrade/downgrade rows. Default: 5.

    Returns:
        dict[str, Any]: A dictionary with the following keys:
            - basic_info: list of key metrics
            - calendar: list of earnings/dividend events (if available)
            - news: list of news headlines with source and date
            - recommendations: analyst recommendations (if available)
            - upgrades_downgrades: recent upgrades/downgrades (if available)

    """
    ticker_obj = yf.Ticker(ticker)

    # Run basic info, calendar, and news in parallel
    with ThreadPoolExecutor() as executor:
        info_future = executor.submit(ticker_obj.info.copy)
        calendar_future = executor.submit(ticker_obj.calendar.copy)
        news_future = executor.submit(ticker_obj.news.copy)

        info = info_future.result()
        if not info:
            raise ValueError(f"No information available for ticker '{ticker}'")

        # Extract essential fields
        essential_fields = {
            'symbol', 'longName', 'currentPrice', 'marketCap', 'volume', 'trailingPE',
            'forwardPE', 'dividendYield', 'beta', 'eps', 'totalRevenue', 'totalDebt',
            'profitMargins', 'operatingMargins', 'returnOnEquity', 'returnOnAssets',
            'revenueGrowth', 'earningsGrowth', 'bookValue', 'priceToBook',
            'enterpriseValue', 'pegRatio', 'trailingEps', 'forwardEps'
        }

        basic_info = []
        for key, value in info.items():
            if key in essential_fields:
                if isinstance(value, (datetime,)):
                    value = value.isoformat()
                basic_info.append({"metric": key, "value": value})

        result: dict[str, Any] = {"basic_info": basic_info}

        # Process calendar (earnings / dividend dates)
        calendar = calendar_future.result()
        if calendar is not None and not calendar.empty:
            calendar_data = []
            for key, val in calendar.items():
                if isinstance(val[0], (datetime,)):
                    val_str = val[0].isoformat()
                else:
                    val_str = str(val[0])
                calendar_data.append({"event": key, "value": val_str})
            result["calendar"] = calendar_data

        # Process news
        news_items = news_future.result()
        if news_items:
            news_data = []
            for item in news_items[:max_news]:
                pub_date = item.get("providerPublishTime")
                date_str = datetime.utcfromtimestamp(pub_date).isoformat() if pub_date else ""
                news_data.append({
                    "date": date_str,
                    "title": item.get("title", "Untitled"),
                    "source": item.get("publisher", "Unknown"),
                    "url": item.get("link", "")
                })
            result["news"] = news_data

    # Run recommendations and upgrades/downgrades in parallel
    with ThreadPoolExecutor() as executor:
        rec_future = executor.submit(ticker_obj.recommendations.copy)
        upgrades_future = executor.submit(ticker_obj.upgrades.copy)

        recommendations = rec_future.result()
        if recommendations is not None and not recommendations.empty:
            rec_csv = recommendations.head(max_recommendations).to_csv(index=True)
            result["recommendations"] = rec_csv

        upgrades = upgrades_future.result()
        if upgrades is not None and not upgrades.empty:
            upgrades_sorted = upgrades.sort_index(ascending=False)
            upgrades_csv = upgrades_sorted.head(max_upgrades).to_csv(index=True)
            result["upgrades_downgrades"] = upgrades_csv

    return result

# 获取财务报表数据（利润表、资产负债表、现金流量表）
@mcp.tool()
def get_financial_statements(
    ticker: str,
    statement_types: list[Literal["income", "balance", "cash"]] = ["income"],
    frequency: Literal["quarterly", "annual"] = "quarterly",
    max_periods: int = 8
) -> dict[str, str]:
    """Get financial statements. Returns dict with statement type as key and CSV data as value."""
    def get_single_statement(stmt_type: str):
        t = yf.Ticker(ticker)
        if stmt_type == "income":
            return t.quarterly_income_stmt if frequency == "quarterly" else t.income_stmt
        elif stmt_type == "balance":
            return t.quarterly_balance_sheet if frequency == "quarterly" else t.balance_sheet
        else:  # cash
            return t.quarterly_cashflow if frequency == "quarterly" else t.cashflow

    # Fetch all requested statements in parallel
    with ThreadPoolExecutor() as executor:
        futures = {stmt_type: executor.submit(get_single_statement, stmt_type) for stmt_type in statement_types}

        for stmt_type, future in futures.items():
            df = future.result()
            if df is None or df.empty:
                raise ValueError(f"No {stmt_type} statement data found for {ticker}")

            if len(df.columns) > max_periods:
                df = df.iloc[:, :max_periods]

            df_reset = df.reset_index()

    return df_reset

# 获取指定股票代码的历史财报数据（收益历史）
@mcp.tool()
def get_earnings_history(ticker: str, max_entries: int = 8) -> str:
    """
    Fetch the historical earnings data of a given stock and return as CSV string.
    
    Args:
        ticker: Stock symbol
        max_entries: Maximum number of entries to return
    
    Returns:
        CSV formatted string of earnings history
    """
    if not ticker or not isinstance(ticker, str):
        raise ValueError("Ticker must be a non-empty string")

    ticker = ticker.upper().strip()
    
    stock = yf.Ticker(ticker)
    earnings_history = stock.earnings_history if hasattr(stock, 'earnings_history') else None
    
    if earnings_history is None or len(earnings_history) == 0:
        try:
            earnings_history = stock.get_earnings_history()
        except Exception:
            earnings_history = None

    if earnings_history is None or len(earnings_history) == 0:
        raise ValueError(f"No earnings history data found for {ticker}")
    
    if not isinstance(earnings_history, pd.DataFrame):
        earnings_history = pd.DataFrame(earnings_history)
    
    earnings_history = earnings_history.head(max_entries)
    
    output = io.StringIO()
    earnings_history.to_csv(output, index=False)
    return output.getvalue()

# 获取指定股票代码的主要机构持股和共同基金持股
@mcp.tool()
def get_institutional_holders(ticker: str, top_n: int = 20) -> dict[str, Any]:
    """
    Fetch major institutional and mutual fund holders of a stock.
    
    Args:
        ticker: Stock symbol
        top_n: Maximum number of top holders to return
    
    Returns:
        A dictionary containing CSV strings for institutional and mutual fund holders
    """
    # Basic ticker validation
    if not ticker or not isinstance(ticker, str):
        raise ValueError("Ticker must be a non-empty string")

    ticker = ticker.upper().strip()
    
    # Function to fetch data from yfinance safely
    def fetch_holders(kind: str) -> pd.DataFrame | None:
        stock = yf.Ticker(ticker)
        try:
            if kind == "institutional":
                data = stock.institutional_holders
            elif kind == "mutualfund":
                data = stock.mutualfund_holders
            else:
                data = None
        except Exception:
            data = None
        # Convert to DataFrame if needed
        if data is not None and not isinstance(data, pd.DataFrame):
            data = pd.DataFrame(data)
        return data

    # Fetch both holders in parallel
    with ThreadPoolExecutor() as executor:
        inst_future = executor.submit(fetch_holders, "institutional")
        fund_future = executor.submit(fetch_holders, "mutualfund")
        inst_holders = inst_future.result()
        fund_holders = fund_future.result()

    # Limit top_n
    if isinstance(inst_holders, pd.DataFrame):
        inst_holders = inst_holders.head(top_n)
    if isinstance(fund_holders, pd.DataFrame):
        fund_holders = fund_holders.head(top_n)

    if (inst_holders is None or inst_holders.empty) and (fund_holders is None or fund_holders.empty):
        raise ValueError(f"No institutional or mutual fund holder data found for {ticker}")

    # Convert to CSV strings
    def df_to_csv(df: pd.DataFrame) -> str:
        output = io.StringIO()
        df.to_csv(output, index=False)
        return output.getvalue()

    result = {"ticker": ticker, "top_n": top_n}
    if inst_holders is not None and not inst_holders.empty:
        result["institutional_holders"] = df_to_csv(inst_holders)
    if fund_holders is not None and not fund_holders.empty:
        result["mutual_fund_holders"] = df_to_csv(fund_holders)

    return result

# 获取指定股票代码的内幕交易数据
@mcp.tool()
def get_insider_trades(ticker: str, max_trades: int = 20) -> str:
    """
    Fetch recent insider trading transactions of a stock and return as CSV string.
    
    Args:
        ticker: Stock symbol
        max_trades: Maximum number of insider trades to return
    
    Returns:
        CSV formatted string of insider trades
    """
    # Basic ticker validation
    if not ticker or not isinstance(ticker, str):
        raise ValueError("Ticker must be a non-empty string")

    ticker = ticker.upper().strip()
    
    # Fetch insider transactions using yfinance
    stock = yf.Ticker(ticker)
    try:
        trades = stock.get_insider_transactions()
    except Exception:
        trades = None

    if trades is None or (isinstance(trades, pd.DataFrame) and trades.empty):
        raise ValueError(f"No insider trading data found for {ticker}")

    # Convert to DataFrame if not already
    if not isinstance(trades, pd.DataFrame):
        trades = pd.DataFrame(trades)
    
    # Limit to max_trades
    trades = trades.head(max_trades)
    
    # Convert DataFrame to CSV string
    output = io.StringIO()
    trades.to_csv(output, index=False)
    return output.getvalue()

# 获取指定股票代码的历史价格数据（OHLCV）
@mcp.tool()
def get_price_history(
    ticker: str,
    period: Literal["1d", "5d", "1mo", "3mo", "6mo", "1y", "2y", "5y", "10y", "ytd", "max"] = "1mo"
) -> str:
    """
    Fetch historical OHLCV (Open, High, Low, Close, Volume) data of a stock and return as CSV string.
    
    Args:
        ticker: Stock symbol
        period: Time period for historical data
    
    Returns:
        CSV formatted string of historical price data
    """
    # Basic ticker validation
    if not ticker or not isinstance(ticker, str):
        raise ValueError("Ticker must be a non-empty string")
    
    ticker = ticker.upper().strip()

    # Determine interval based on period
    interval = "1mo" if period in ["2y", "5y", "10y", "max"] else "1d"

    # Fetch historical data using yfinance
    stock = yf.Ticker(ticker)
    try:
        history = stock.history(period=period, interval=interval)
    except Exception:
        history = None

    if history is None or history.empty:
        raise ValueError(f"No historical data found for {ticker}")

    # Reset index to include date as a column
    history_with_dates = history.reset_index()
    history_with_dates['Date'] = pd.to_datetime(history_with_dates['Date']).dt.strftime('%Y-%m-%d')

    # Convert DataFrame to CSV string
    output = io.StringIO()
    history_with_dates.to_csv(output, index=False)
    return output.getvalue()

# 获取指定股票代码的期权数据
@mcp.tool()
def get_options(
    ticker_symbol: str,
    num_options: int = 10,
    start_date: str | None = None,
    end_date: str | None = None,
    strike_lower: float | None = None,
    strike_upper: float | None = None,
    option_type: Literal["C", "P"] | None = None,
) -> str:
    """
    Fetch options data for a stock and return as CSV string.
    
    Args:
        ticker_symbol: Stock symbol
        num_options: Maximum number of options rows to return
        start_date: Filter options expiring on or after this date (YYYY-MM-DD)
        end_date: Filter options expiring on or before this date (YYYY-MM-DD)
        strike_lower: Minimum strike price
        strike_upper: Maximum strike price
        option_type: 'C' for calls, 'P' for puts, None for both
    
    Returns:
        CSV formatted string of options data
    """
    # Basic ticker validation
    if not ticker_symbol or not isinstance(ticker_symbol, str):
        raise ValueError("Ticker must be a non-empty string")
    ticker_symbol = ticker_symbol.upper().strip()

    # Parse dates
    if start_date:
        try:
            start_date_parsed = pd.to_datetime(start_date)
        except Exception:
            raise ValueError("Invalid start_date format, should be YYYY-MM-DD")
    else:
        start_date_parsed = None

    if end_date:
        try:
            end_date_parsed = pd.to_datetime(end_date)
        except Exception:
            raise ValueError("Invalid end_date format, should be YYYY-MM-DD")
    else:
        end_date_parsed = None

    # Fetch options expirations
    stock = yf.Ticker(ticker_symbol)
    expirations = stock.options
    if not expirations:
        raise ValueError(f"No options available for {ticker_symbol}")

    # Filter expirations by date range
    valid_expirations = [
        exp for exp in expirations
        if ((not start_date_parsed or pd.to_datetime(exp) >= start_date_parsed) and
            (not end_date_parsed or pd.to_datetime(exp) <= end_date_parsed))
    ]
    if not valid_expirations:
        raise ValueError(f"No options found for {ticker_symbol} within specified date range")

    # Function to fetch options chain safely
    def fetch_chain(expiry: str) -> pd.DataFrame | None:
        try:
            chain = stock.option_chain(expiry)
            df = pd.DataFrame()
            if option_type == "C":
                df = chain.calls
            elif option_type == "P":
                df = chain.puts
            else:
                df = pd.concat([chain.calls, chain.puts], ignore_index=True)
            df['expiryDate'] = expiry
            return df
        except Exception:
            return None

    # Fetch all valid chains in parallel
    with ThreadPoolExecutor() as executor:
        chains = list(filter(None, executor.map(fetch_chain, valid_expirations)))

    if not chains:
        raise ValueError(f"No options found for {ticker_symbol} matching criteria")

    # Concatenate all chains
    df_all = pd.concat(chains, ignore_index=True)

    # Apply strike filters
    if strike_lower is not None:
        df_all = df_all[df_all['strike'] >= strike_lower]
    if strike_upper is not None:
        df_all = df_all[df_all['strike'] <= strike_upper]

    # Sort by openInterest and volume descending
    df_all = df_all.sort_values(['openInterest', 'volume'], ascending=[False, False])

    # Take top N options
    df_subset = df_all.head(num_options)

    # Convert to CSV string
    output = io.StringIO()
    df_subset.to_csv(output, index=False)
    return output.getvalue()


GET_TICKER_DATA_SCHEMA = {
    "type": "function",
    "function": {
        "name": "get_ticker_data",
        "description": "Get comprehensive data for a given stock ticker, including basic metrics, earnings calendar, news headlines, analyst recommendations, and upgrades/downgrades history.",
        "parameters": {
            "type": "object",
            "properties": {
                "ticker": {
                    "type": "string",
                    "description": "The stock ticker symbol (e.g., 'AAPL')."
                },
                "max_news": {
                    "type": "integer",
                    "description": "Maximum number of news headlines to return. Default is 5."
                },
                "max_recommendations": {
                    "type": "integer",
                    "description": "Maximum number of analyst recommendation rows to return. Default is 5."
                },
                "max_upgrades": {
                    "type": "integer",
                    "description": "Maximum number of upgrade/downgrade rows to return. Default is 5."
                }
            },
            "required": ["ticker"]
        }
    }
}

GET_FINANCIAL_STATEMENTS_SCHEMA = {
    "type": "function",
    "function": {
        "name": "get_financial_statements",
        "description": "Retrieve financial statements (income statement, balance sheet, cash flow) for a given stock ticker.",
        "parameters": {
            "type": "object",
            "properties": {
                "ticker": {
                    "type": "string",
                    "description": "The stock ticker symbol (e.g., 'MSFT')."
                },
                "statement_types": {
                    "type": "array",
                    "items": {
                        "type": "string",
                        "enum": ["income", "balance", "cash"]
                    },
                    "description": "List of statement types to retrieve. Options: 'income', 'balance', 'cash'. Default is ['income']."
                },
                "frequency": {
                    "type": "string",
                    "enum": ["quarterly", "annual"],
                    "description": "Frequency of financial statements. Either 'quarterly' or 'annual'. Default is 'quarterly'."
                },
                "max_periods": {
                    "type": "integer",
                    "description": "Maximum number of reporting periods to return. Default is 8."
                }
            },
            "required": ["ticker"]
        }
    }
}

GET_EARNINGS_HISTORY_SCHEMA = {
    "type": "function",
    "function": {
        "name": "get_earnings_history",
        "description": "Fetch historical earnings data for a given stock ticker and return it as a CSV string.",
        "parameters": {
            "type": "object",
            "properties": {
                "ticker": {
                    "type": "string",
                    "description": "The stock ticker symbol (e.g., 'GOOG')."
                },
                "max_entries": {
                    "type": "integer",
                    "description": "Maximum number of historical earnings entries to return. Default is 8."
                }
            },
            "required": ["ticker"]
        }
    }
}

GET_INSTITUTIONAL_HOLDERS_SCHEMA = {
    "type": "function",
    "function": {
        "name": "get_institutional_holders",
        "description": "Retrieve major institutional and mutual fund holders for a given stock ticker, returning data as CSV strings.",
        "parameters": {
            "type": "object",
            "properties": {
                "ticker": {
                    "type": "string",
                    "description": "The stock ticker symbol (e.g., 'TSLA')."
                },
                "top_n": {
                    "type": "integer",
                    "description": "Maximum number of top holders to return for each category. Default is 20."
                }
            },
            "required": ["ticker"]
        }
    }
}

GET_INSIDER_TRADES_SCHEMA = {
    "type": "function",
    "function": {
        "name": "get_insider_trades",
        "description": "Fetch recent insider trading transactions for a given stock ticker and return as a CSV string.",
        "parameters": {
            "type": "object",
            "properties": {
                "ticker": {
                    "type": "string",
                    "description": "The stock ticker symbol (e.g., 'AMZN')."
                },
                "max_trades": {
                    "type": "integer",
                    "description": "Maximum number of insider trades to return. Default is 20."
                }
            },
            "required": ["ticker"]
        }
    }
}

GET_PRICE_HISTORY_SCHEMA = {
    "type": "function",
    "function": {
        "name": "get_price_history",
        "description": "Fetch historical OHLCV (Open, High, Low, Close, Volume) price data for a given stock ticker and return as a CSV string.",
        "parameters": {
            "type": "object",
            "properties": {
                "ticker": {
                    "type": "string",
                    "description": "The stock ticker symbol (e.g., 'NVDA')."
                },
                "period": {
                    "type": "string",
                    "enum": ["1d", "5d", "1mo", "3mo", "6mo", "1y", "2y", "5y", "10y", "ytd", "max"],
                    "description": "The time period for historical data. Default is '1mo'."
                }
            },
            "required": ["ticker"]
        }
    }
}

GET_OPTIONS_SCHEMA = {
    "type": "function",
    "function": {
        "name": "get_options",
        "description": "Fetch options chain data for a given stock ticker, with optional filters for date range, strike price range, and option type. Returns data as a CSV string.",
        "parameters": {
            "type": "object",
            "properties": {
                "ticker_symbol": {
                    "type": "string",
                    "description": "The stock ticker symbol (e.g., 'AAPL')."
                },
                "num_options": {
                    "type": "integer",
                    "description": "Maximum number of options rows to return. Default is 10."
                },
                "start_date": {
                    "type": "string",
                    "description": "Filter for options expiring on or after this date (YYYY-MM-DD)."
                },
                "end_date": {
                    "type": "string",
                    "description": "Filter for options expiring on or before this date (YYYY-MM-DD)."
                },
                "strike_lower": {
                    "type": "number",
                    "description": "Minimum strike price filter."
                },
                "strike_upper": {
                    "type": "number",
                    "description": "Maximum strike price filter."
                },
                "option_type": {
                    "type": "string",
                    "enum": ["C", "P"],
                    "description": "Option type: 'C' for calls, 'P' for puts. If not provided, both are returned."
                }
            },
            "required": ["ticker_symbol"]
        }
    }
}


async def main():
    client = Client(mcp)
    async with client:
        tools = await client.list_tools()
        print("Available tools:", tools)

        result = await client.call_tool("get_ticker_data", {"ticker": "AAPL", "max_news": 3})
        print("Market Movers Result:", result)

        result = await client.call_tool("get_financial_statements", {
            "ticker":"AAPL",
            "statement_types":["income", "balance"],
            "frequency":"annual",
            "max_periods":4
        })
        print("Market Movers Result:", result)

        result = await client.call_tool(
            "get_earnings_history",{
                "ticker": "AAPL",
                "max_entries": 5
        })
        print("Market Movers Result:", result)

        result = await client.call_tool(
        "get_institutional_holders",{
            "ticker": "AAPL",
            "top_n": 5
        })
        print("Market Movers Result:", result)

        result = await client.call_tool(
        "get_insider_trades",{
            "ticker": "AAPL",
            "max_trades": 5
        })
        print("Market Movers Result:", result)

        result = await client.call_tool(
        "get_price_history",{
            "ticker": "AAPL",
            "period": "3mo"
        })
        print("Market Movers Result:", result)

        result = await client.call_tool(
        "get_options",{
            "ticker_symbol": "AAPL",
            "num_options": 5,
            "option_type": "C"
        })
        print("Market Movers Result:", result)


if __name__ == "__main__":
    asyncio.run(main())